package com.etl.dataflow.executor.core;

import cn.hutool.core.bean.BeanUtil;
import com.etl.dataflow.common.entity.JdbcProperties;
import com.etl.dataflow.common.entity.TaskParam;
import com.etl.dataflow.common.entity.TaskStatisticInfo;
import com.etl.dataflow.common.enums.TaskStateEnum;
import com.etl.dataflow.executor.element.Row;
import com.etl.dataflow.executor.entity.ExecutorTaskInfo;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author dx
 * @since 2022/4/20
 */
@SuppressWarnings("AlibabaThreadPoolCreation")
public class TaskRunner {
    private final ExecutorTaskInfo executorTaskInfo;
    private final TaskParam taskParam;
    private final TaskContext taskContext;
    private final AbstractTaskLaunchInterceptor taskLaunchInterceptor;
    private final AbstractAfterReadInterceptor afterReadInterceptor;
    private final AbstractAfterWriteInterceptor afterWriteInterceptor;
    private volatile boolean isShutdown = false;
    private Reader mReader;
    private Writer mWriter;
    private OnTaskFinishedListener onTaskFinishedListener;
    private OnTaskErrorListener onTaskErrorListener;
    private long taskLaunchTime = 0L;

    public TaskRunner(ExecutorTaskInfo taskInfo) {
        this.executorTaskInfo = taskInfo;
        this.taskParam = taskInfo.getTaskParam();
        this.taskContext = new TaskContext().setTaskInfo(taskInfo);
        this.taskLaunchInterceptor = taskInfo.getTaskLaunchInterceptor();
        this.afterReadInterceptor = taskInfo.getAfterReadInterceptor();
        this.afterWriteInterceptor = taskInfo.getAfterWriteInterceptor();
    }

    public void setOnTaskFinishedListener(OnTaskFinishedListener onTaskFinishedListener) {
        this.onTaskFinishedListener = onTaskFinishedListener;
    }

    public void setOnTaskErrorListener(OnTaskErrorListener onTaskErrorListener) {
        this.onTaskErrorListener = onTaskErrorListener;
    }

    public void run() {
        taskLaunchTime = System.currentTimeMillis();
        try {
            this.taskContext.getTaskInfo().setState(TaskStateEnum.RUNNING);
            launchInterceptorAndReaderAndWriter();
        } catch (Throwable e) {
            this.taskContext.getTaskInfo().setState(TaskStateEnum.FATAL_ERROR);
            if (onTaskErrorListener != null) {
                onTaskErrorListener.onError(new Exception(e));
            }
            throw e;
        }
    }

    @SuppressWarnings("AlibabaAvoidManuallyCreateThread")
    private void launchInterceptorAndReaderAndWriter() {
        // 执行任务启动拦截器
        invokeTaskLaunchInterceptor();

        JdbcProperties readerProperties = taskParam.getSrc();
        JdbcProperties writerProperties = taskParam.getDst();
        MemoryChannel<Row> memoryChannel = new MemoryChannel<>(taskParam.getBufferSize());
        Transporter transporter = new Transporter(memoryChannel);

        new Thread(() -> {
            // 将taskContext写入InheritableThreadLocal
            TaskContextHolder.setTaskContext(this.taskContext);

            RdbmsReader reader = new OracleReader(taskParam.getSrcTableName(), readerProperties, transporter);
            reader.setSql(taskParam.getSelectSql());
            reader.setAfterReadInterceptor(this.afterReadInterceptor);

            reader.setOnReadErrorListener(e -> {
                if (onTaskErrorListener != null) {
                    onTaskErrorListener.onError(e);
                }
            });
            this.mReader = reader;
            reader.startRead();
        }, "reader-thread-" + taskParam.getDstTableName()).start();

        new Thread(() -> {
            // 将taskContext写入InheritableThreadLocal
            TaskContextHolder.setTaskContext(this.taskContext);
            WriterFactory writerFactory = (WriterFactory) taskParam.getExtensionMap().get("writerFactory");
            if (writerFactory == null) {
                writerFactory = new OracleInsertOnlyWriterFactory();
            }
            Writer writer = writerFactory.getWriter(taskParam.getDstTableName(), writerProperties, transporter);
            writer.setAfterWriteInterceptor(this.afterWriteInterceptor);

            writer.setOnWriteFinishedListener(() -> {
                if (onTaskFinishedListener != null) {
                    this.taskContext.getTaskInfo().setState(TaskStateEnum.FINISHED);
                    onTaskFinishedListener.onFinished();
                }
            });
            writer.setOnFatalErrorListener(e -> {
                if (onTaskErrorListener != null) {
                    this.taskContext.getTaskInfo().setState(TaskStateEnum.FATAL_ERROR);
                    onTaskErrorListener.onError(e);
                }
            });
            this.mWriter = writer;
            writer.startWrite();
        }, "writer-thread-" + taskParam.getDstTableName()).start();
    }

    private void invokeTaskLaunchInterceptor() {
        if (taskLaunchInterceptor != null) {
            TaskParam copyOfTaskParam = new TaskParam();
            BeanUtil.copyProperties(this.taskParam, copyOfTaskParam);
            // 启动拦截器可能会修改任务参数, 这里需要将修改后的参数写回原本的参数实例
            TaskParam newTaskParam = taskLaunchInterceptor.beforeLaunch(copyOfTaskParam);
            if (newTaskParam != null) {
                BeanUtil.copyProperties(newTaskParam, taskParam);
            }
            Map<String, Object> extensionMap = taskParam.getExtensionMap();
            if (extensionMap == null) {
                extensionMap = new ConcurrentHashMap<>(8);
                taskParam.setExtensionMap(extensionMap);
            }
            if (extensionMap.get("writerFactory") == null) {
                extensionMap.put("writerFactory", new OracleInsertOnlyWriterFactory());
            }
        }
    }

    public TaskStatisticInfo getTaskStatisticInfo() {
        return this.taskContext.getTaskStatisticInfo();
    }

    public synchronized void shutdown() {
        if (isShutdown) {
            return;
        }
        if (mReader != null) {
            mReader.exit();
        }
        if (mWriter != null) {
            mWriter.exit();
        }
        this.isShutdown = true;
    }

    public ExecutorTaskInfo getExecutorTaskInfo() {
        return executorTaskInfo;
    }

    public long getTaskLaunchTime() {
        return taskLaunchTime;
    }

    public Writer getWriter() {
        return this.mWriter;
    }

    public Reader getReader() {
        return this.mReader;
    }

    public boolean isShutdown() {
        return this.isShutdown;
    }

    public interface OnTaskFinishedListener {
        /**
         * 任务正常结束回调
         */
        void onFinished();
    }

    public interface OnTaskErrorListener {
        /**
         * 任务错误回调
         *
         * @param e Exception
         */
        void onError(Exception e);
    }

}
